import logging
import os
import numpy as np

from control_models.base import ALLOW_CUSTOM_MODEL_PATH, MODEL_ROOT, MODEL_SCORE_THRESHOLD, ControlModel,_model_cache, get_bool
from typing import ClassVar, List, Dict
from paddleocr import PaddleOCR

from label_studio_ml.model import LabelStudioMLBase
from label_studio_ml.utils import DATA_UNDEFINED_NAME
from label_studio_sdk.label_interface.control_tags import ControlTag
from label_studio_sdk.label_interface import LabelInterface


logger = logging.getLogger(__name__)


class PaddleOCRModel(ControlModel):
    """
    Class representing a Choices (classes) control tag for YOLO model.
    """

    type = "Textarea"
    textline_orientation_model_name: ClassVar[str]= ""
    text_detection_model_name: ClassVar[str] = "PP-OCRv5_server_det" # 文本检测模型
    text_recognition_model_name: ClassVar[str] = "PP-OCRv5_server_rec" # 文本识别模型
    model: PaddleOCR

    @classmethod
    def is_control_matched(cls, control,mode=None) -> bool:
        # check object tag type
        if control.objects[0].tag != "Image":
            return False
        # support both Choices and Taxonomy because of their similarity
        return control.tag in [cls.type, "Taxonomy"] and mode is None

    def predict_ocr(self, image) -> List[Dict]:
        results = self.model.predict(image)
        # self.debug_plot(results[0].plot())
        return self.create_textarea(results, image)

    def create_textarea(self, results, path):
        logger.debug(f"create_choices: {self.from_name}")
        mode = self.control.attr.get("choice", "single")
        result = results[0]
        score = np.mean(result["rec_scores"])

        rec_texts = "".join(result["rec_texts"])
        logger.debug(
            "----------------------\n"
            f"task id > {path}\n"
            f"control: {self.control}\n"
            f"result > {result}\n"
            f"score > {score}\n"
            f"rec_texts > {rec_texts}\n"
        )

        if score < self.model_score_threshold:
            logger.debug(f"Score is too low for single choice: {names[0]} = {probs[0]}")
            return []
        # add new region with rectangle
        return {
                "from_name": self.from_name,
                "to_name": self.to_name,
                "type": "textarea",
                "rec_texts": rec_texts,
                "rec_scores": float(score),
        }


    @classmethod
    def load_paddleocr_model(cls, text_detection_model_name,text_recognition_model_name) -> PaddleOCR:
        """Load paddleocr model from the file."""
        text_detection_model_dir = os.path.join(MODEL_ROOT, text_detection_model_name)
        text_recognition_model_dir = os.path.join(MODEL_ROOT, text_recognition_model_name)
        logger.info(f"Loading paddleocr detection model: {text_detection_model_dir} {os.path.exists(text_detection_model_dir)}")
        logger.info(f"Loading paddleocr recognition model: {text_recognition_model_dir} {os.path.exists(text_recognition_model_dir)}")
        model = PaddleOCR(
            text_detection_model_name="PP-OCRv5_server_det",
            text_recognition_model_name="PP-OCRv5_server_rec",
            text_detection_model_dir=text_detection_model_dir,
            text_recognition_model_dir=text_recognition_model_dir,
            use_doc_orientation_classify=False, 
            use_doc_unwarping=False, 
            use_textline_orientation=False
        ) 
        logger.info(f"Model {text_detection_model_dir} names:\n{model}")
        return model

    @classmethod
    def get_cached_model(cls, text_detection_model_name,text_recognition_model_name) -> PaddleOCR:
        path = f"{text_detection_model_name},{text_recognition_model_name}" # 
        if path not in _model_cache:
            _model_cache[path] = cls.load_paddleocr_model(text_detection_model_name,text_recognition_model_name)
        return _model_cache[path]
    
    
    @classmethod
    def create(cls, mlbackend: LabelStudioMLBase, control: ControlTag):
        """Factory method to create an instance of a specific control model class.
        Args:
            mlbackend (LabelStudioMLBase): The ML backend instance.
            control (ControlTag): The control tag from the Label Studio Interface.
        """
        from_name = control.name
        to_name = control.to_name[0]
        value = control.objects[0].value_name

        # 判断模型是否需要ocr
        if not get_bool(control.attr, "model_ocr", "false"):
            logger.info(
                f"Skipping control tag '{control.tag}' with name '{from_name}', model_ocr=false found"
            )
            return None
        # read threshold attribute from the control tag, e.g.: <RectangleLabels model_score_threshold="0.5">
        model_score_threshold = float(
            control.attr.get("model_score_threshold")
            or control.attr.get(
                "score_threshold"
            )  # not recommended option, use `model_score_threshold`
            or MODEL_SCORE_THRESHOLD
        )
        # read `model_path` attribute from the control tag
        text_detection_model_name = (
            ALLOW_CUSTOM_MODEL_PATH and control.attr.get("text_detection_model_name")
        ) or cls.text_detection_model_name
        text_recognition_model_name = (
            ALLOW_CUSTOM_MODEL_PATH and control.attr.get("text_recognition_model_name")
        ) or cls.text_recognition_model_name

        model = cls.get_cached_model(text_detection_model_name=text_detection_model_name,
                                     text_recognition_model_name=text_recognition_model_name)

        return cls(
            control=control,
            from_name=from_name,
            to_name=to_name,
            value=value,
            model=model,
            model_score_threshold=model_score_threshold,
            label_studio_ml_backend=mlbackend,
            project_id=mlbackend.project_id,
        )
# pre-load and cache default model at startup
# PaddleOCRModel.get_cached_model(PaddleOCRModel.text_detection_model_name,
#                                 PaddleOCRModel.text_recognition_model_name)
